# scepter_motion_final.py - 最终修复版本
import time
import os

# 设置关键环境变量
os.environ['ROS_PACKAGE_PATH'] = '/home/parallels/PycharmProjects/aghanim-scepter/src:/opt/ros/jazzy/share'

from moveit.task_constructor import core, stages
import rclcpp


def main():
    rclcpp.init()
    node = rclcpp.Node("scepter_mtc_final")
    
    print("🚀 Starting Scepter MTC Demo...")
    
    # 设置规划组
    arm_group = "arm_group"
    gripper_group = "gripper_group"

    # 创建任务并加载机器人模型
    task = core.Task()
    task.name = "scepter_final_demo"
    
    try:
        task.loadRobotModel(node)
        print("✅ Robot model loaded successfully")
    except Exception as e:
        print(f"❌ Failed to load robot model: {e}")
        return False

    # 创建关节插值规划器
    joint_planner = core.JointInterpolationPlanner()
    print("✅ Joint space planner created")
    
    # 从当前状态开始
    current_state = stages.CurrentState("current state")
    task.add(current_state)

    print("📋 Building task sequence...")

    # 任务序列：演示基本的机械臂运动
    stages_info = []
    
    # 1. 移动到准备位置
    move_to_ready = stages.MoveTo("move to ready", joint_planner)
    move_to_ready.group = arm_group
    move_to_ready.setGoal("ready_pose")
    task.add(move_to_ready)
    stages_info.append("Move to ready position")

    # 2. 操作夹爪
    open_gripper = stages.MoveTo("open gripper", joint_planner)
    open_gripper.group = gripper_group
    open_gripper.setGoal("open_pose")
    task.add(open_gripper)
    stages_info.append("Open gripper")

    # 3. 移动到完成位置
    move_to_finish = stages.MoveTo("move to finish", joint_planner)
    move_to_finish.group = arm_group
    move_to_finish.setGoal("finish_pose")
    task.add(move_to_finish)
    stages_info.append("Move to finish position")

    # 4. 关闭夹爪
    close_gripper = stages.MoveTo("close gripper", joint_planner)
    close_gripper.group = gripper_group
    close_gripper.setGoal("close_pose")
    task.add(close_gripper)
    stages_info.append("Close gripper")

    # 5. 回到初始位置
    move_home = stages.MoveTo("move home", joint_planner)
    move_home.group = arm_group
    move_home.setGoal("stand_pose")
    task.add(move_home)
    stages_info.append("Return to home position")

    print(f"📊 Task created with {len(stages_info)} stages")
    
    # 执行规划
    print("🔄 Planning trajectory...")
    try:
        planning_success = task.plan()
        
        if planning_success:
            print("✅ Planning succeeded!")
            
            # 检查解决方案
            if hasattr(task, 'solutions') and len(task.solutions) > 0:
                solution = task.solutions[0]
                print(f"🎯 Found {len(task.solutions)} solution(s)")
                print("📤 Publishing solution...")
                task.publish(solution)
                print("✅ Solution published successfully")
                
                # 显示阶段信息
                print("📋 Executed stages:")
                for i, stage_name in enumerate(stages_info):
                    print(f"  {i+1}. {stage_name}")
                    
                return True
            else:
                print("⚠️ Planning succeeded but no solutions found")
                return False
        else:
            print("❌ Planning failed")
            print("💡 This could be due to:")
            print("  - Invalid target poses")
            print("  - Collision constraints")
            print("  - Joint limits")
            print("  - Missing kinematics configuration")
            return False
            
    except Exception as e:
        print(f"❌ Error during planning: {e}")
        print(f"❌ Error type: {type(e).__name__}")
        return False

    finally:
        print("🏁 Demo execution completed")


def print_system_info():
    """打印系统配置信息用于调试"""
    print("\n🔍 System Configuration:")
    print(f"  ROS_PACKAGE_PATH: {os.environ.get('ROS_PACKAGE_PATH', 'Not set')}")
    print(f"  Current working directory: {os.getcwd()}")
    

if __name__ == "__main__":
    try:
        print_system_info()
        success = main()
        
        if success:
            print("\n🎉 Demo completed successfully!")
            print("⏳ Waiting 3 seconds...")
            time.sleep(3)
        else:
            print("\n⚠️ Demo completed with issues")
            
    except KeyboardInterrupt:
        print("\n🛑 Demo interrupted by user")
    except Exception as e:
        print(f"\n💥 Unexpected error: {e}")
        print("🔧 Please check your MoveIt configuration and try again")